package com.hyk.learning.registry.zk;

import com.hyk.learning.common.enums.LoadBalanceEnum;
import com.hyk.learning.common.enums.RpcErrorMessageEnum;
import com.hyk.learning.common.exception.RpcException;
import com.hyk.learning.common.extension.ExtensionLoader;
import com.hyk.learning.common.utils.CollectionUtil;
import com.hyk.learning.dto.RpcRequest;
import com.hyk.learning.loadbalance.LoadBalance;
import com.hyk.learning.registry.ServiceDiscovery;
import com.hyk.learning.registry.zk.util.CuratorUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.curator.framework.CuratorFramework;
import java.net.InetSocketAddress;
import java.util.List;

/**
 * @className: ZkServiceDiscoveryImpl
 * @author: hyk199710
 * @description: service discovery based on zookeeper
 * @date: 2024/4/29 22:56
 */
@Slf4j
public class ZkServiceDiscoveryImpl implements ServiceDiscovery {

    private final LoadBalance loadBalance;
    private static final String RETRY = "canRetry";

    public ZkServiceDiscoveryImpl() throws Exception {
        this.loadBalance = ExtensionLoader.getExtensionLoader(LoadBalance.class)
                .getExtension(LoadBalanceEnum.LOADBALANCE.getName());
    }

    @Override
    public InetSocketAddress lookupService(RpcRequest rpcRequest) {
        String rpcServiceName = rpcRequest.getRpcServiceName();
        CuratorFramework zkClient = CuratorUtils.getZkClient();
        List<String> serviceUrlList = CuratorUtils.getChildrenNodes(zkClient, rpcServiceName);
        if (CollectionUtil.isEmpty(serviceUrlList)) {
            throw new RpcException(RpcErrorMessageEnum.SERVICE_CAN_NOT_BE_FOUND, rpcServiceName);
        }
        //负载均衡策略
        String targetServiceUrl = loadBalance.selectServiceAddress(serviceUrlList, rpcRequest);
        log.info("Successfully found the service address:[{}]", targetServiceUrl);
        String[] socketAddressArray = targetServiceUrl.split(":");
        String host = socketAddressArray[0];
        int port = Integer.parseInt(socketAddressArray[1]);
        return new InetSocketAddress(host, port);
    }

    @Override
    public boolean checkRetry(String serviceName) {
        boolean canRetry = false;
        try {
            CuratorFramework zkClient = CuratorUtils.getZkClient();
            List<String> serviceList = zkClient.getChildren().forPath("/" + RETRY);
            for(String s : serviceList){
                if(s.equals(serviceName)){
                    log.info("Service {} was found in white list!", serviceName);
                    canRetry = true;
                    break;
                }
            }
        } catch (Exception e) {
            log.error("Retry Service search failed!");
            throw new RuntimeException(e);
        }

        return canRetry;
    }
}
